package middleware.optimize;

import middleware.llvm.BasicBlock;
import middleware.llvm.Module;
import middleware.llvm.instruction.AllocaInstruction;
import middleware.llvm.instruction.LoadInstruction;
import middleware.llvm.instruction.StoreInstruction;
import middleware.llvm.value.FunctionIRType;

import java.util.HashSet;
import java.util.Map;
import java.util.Set;

/**
 * 死存储优化
 * @author 杨凯伟
 */
public class DeadStoreElimination {
    private final Module irModule;
    private Map<BasicBlock, Set<AllocaInstruction>> outSets;
    private Set<AllocaInstruction> varAllocaInstSet;

    public DeadStoreElimination(Module irModule) {
        this.irModule = irModule;
    }

    public Module pass() {
        irModule.getFunctions().forEach(this::passFunc);
        return irModule;
    }

    private void passFunc(FunctionIRType func) {
        var analyzer = new LiveVariableAnalyzer(func);
        analyzer.analyze();
        outSets = analyzer.getOutSets();

        varAllocaInstSet = new HashSet<>(func.getFirstBasicBlock().getInstructions()
                .stream()
                .filter(inst -> inst instanceof AllocaInstruction)
                .map(inst -> (AllocaInstruction) inst)
                .filter(allocaInstruction -> {
                    var dataType = allocaInstruction.getDataType();
                    return dataType.getPtrNum() == 0 && dataType.getArrayDims().isEmpty();
                })
                .toList());


        func.getBasicBlocks().forEach(this::passBlock);
    }

    private void passBlock(BasicBlock block) {
        var outSet = outSets.get(block);
        Set<AllocaInstruction> needStore = new HashSet<>(outSet);

        var insts = block.getInstructions();
        for (int i = insts.size() - 1; i >= 0; i--) {
            var inst = insts.get(i);
            if (inst instanceof LoadInstruction loadInstruction) {
                var ptr = loadInstruction.getPtr();
                if (ptr instanceof AllocaInstruction allocaInstPtr && varAllocaInstSet.contains(allocaInstPtr)) {
                    needStore.add(allocaInstPtr);
                }
            } else if (inst instanceof StoreInstruction storeInstruction) {
                var ptr = storeInstruction.getPtr();
                if (ptr instanceof AllocaInstruction allocaInstPtr && varAllocaInstSet.contains(allocaInstPtr)) {
                    if (needStore.contains(allocaInstPtr)) {
                        needStore.remove(allocaInstPtr);
                    } else {
                        inst.replaceAllUseWith(null, false); // useList of store inst is empty, so null is safe
                    }
                }
            }
        }
    }
}
